//
// Created by kai-xu on 2022/12/16.
//
#include "preprocess/pro_mat_utils.h"

__global__ void circle_ray_cal_idx_kernel(
    circle_ray::custom_coor_point* key_value, float* aug_mat, float* ego_mat,
    int32_t n_cam, int32_t dz, int32_t width, int32_t height,
    const float voxel_x, const float voxel_y, const float voxel_z,
    const float coors_x_min, const float coors_y_min, const float coors_z_min,
    const float par_x, const float par_y) {
  auto j = std::int32_t(blockIdx.x * blockDim.x + threadIdx.x);  // width
  auto i = std::int32_t(blockIdx.y * blockDim.y + threadIdx.y);  // height
  auto k = std::int32_t(blockIdx.z * blockDim.z + threadIdx.z);  // d * n_cam

  if ((i < height) && (j < width) && (k < n_cam * dz)) {
    auto mat_id = int32_t(k / dz);
    auto dz_index = float(int32_t(k % dz) + 1);
    float i_ = float(i) * par_y;
    float j_ = float(j) * par_x;

    float* aug_mat_add = aug_mat + mat_id * 16;
    float d2 = aug_mat_add[8] * j_ + aug_mat_add[9] * i_ +
               aug_mat_add[10] * dz_index + aug_mat_add[11];
    float w2 = (aug_mat_add[0] * j_ + aug_mat_add[1] * i_ +
                aug_mat_add[2] * dz_index + aug_mat_add[3] * 1) *
               d2;
    float h2 = (aug_mat_add[4] * j_ + aug_mat_add[5] * i_ +
                aug_mat_add[6] * dz_index + aug_mat_add[7] * 1) *
               d2;

    float* ego_mat_add = ego_mat + mat_id * 16;
    float d3 = ego_mat_add[8] * w2 + ego_mat_add[9] * h2 +
               ego_mat_add[10] * d2 + ego_mat_add[11];
    float w3 = ego_mat_add[0] * w2 + ego_mat_add[1] * h2 + ego_mat_add[2] * d2 +
               ego_mat_add[3];
    float h3 = ego_mat_add[4] * w2 + ego_mat_add[5] * h2 + ego_mat_add[6] * d2 +
               ego_mat_add[7];
    auto coors_x = (w3 - coors_x_min) / voxel_x;
    auto coors_y = (h3 - coors_y_min) / voxel_y;
    auto coors_z = (d3 - coors_z_min) / voxel_z;

    key_value[k * width * height + i * width + j].xf = coors_x;
    key_value[k * width * height + i * width + j].yf = coors_y;
    key_value[k * width * height + i * width + j].zf = coors_z;
  }
}

__global__ void circle_ray_reduce_hdim(
    circle_ray::custom_coor_point* o_coor_value,
    const circle_ray::custom_coor_point* i_coor_value, int32_t bs, int32_t nc,
    int32_t dz, int32_t height, int32_t width) {
  auto j = std::int32_t(blockIdx.x * blockDim.x + threadIdx.x);  // width
  auto i = std::int32_t(blockIdx.y * blockDim.y + threadIdx.y);  //  bs * nc *dz
  // data format:[bs, nc, dz, height, width, 3]
  if ((j < width) && (i < bs * nc * dz)) {
    for (int32_t k{0}; k < height; ++k) {
      o_coor_value[(i * width + j)].xf +=
          i_coor_value[(i * height * width + k * width + j)].xf;
      o_coor_value[(i * width + j)].yf +=
          i_coor_value[(i * height * width + k * width + j)].yf;
      o_coor_value[(i * width + j)].zf +=
          i_coor_value[(i * height * width + k * width + j)].zf;
    }
    o_coor_value[(i * width + j)].x =
        int32_t(o_coor_value[(i * width + j)].xf / float(height));
    o_coor_value[(i * width + j)].y =
        int32_t(o_coor_value[(i * width + j)].yf / float(height));
    o_coor_value[(i * width + j)].z =
        int32_t(o_coor_value[(i * width + j)].zf / float(height));
  }
}

__global__ void permuteData_kernel(
    circle_ray::custom_coor_point* o_coor_value,
    const circle_ray::custom_coor_point* i_coor_value, int32_t* idx_ptr,
    const int32_t nthreads, const int32_t bs, const int32_t nc,
    const int32_t dz, const int32_t width, const int32_t out_height,
    const int32_t out_width) {
  // data format:[bs, nc, dz, width, 3] -> [bs, nc , width , dz, 3]
  auto index = std::int32_t(blockIdx.x * blockDim.x + threadIdx.x);
  if (index < nthreads) {
    const int w = index % width;
    const int d = (index / width) % dz;
    const int c = (index / width / dz) % nc;
    const int b = index / width / dz / nc;
    const int new_index = b * nc * dz * width + c * dz * width + w * dz + d;
    if (i_coor_value[index].x < 0 || i_coor_value[index].x >= out_width ||
        i_coor_value[index].y < 0 || i_coor_value[index].y >= out_height) {
      // invalid points
      o_coor_value[new_index].x = out_width / 2;
      o_coor_value[new_index].y = out_height / 2;
    } else {
      o_coor_value[new_index].x = i_coor_value[index].x;
      o_coor_value[new_index].y = i_coor_value[index].y;
      o_coor_value[new_index].z = i_coor_value[index].z;
    }
    idx_ptr[new_index] =
        o_coor_value[new_index].y * out_width + o_coor_value[new_index].x;
  }
}
__global__ void get_circle_map_kernel(float* circle_map, const int32_t* idx_ptr,
                                      int32_t bs, int32_t dz, int32_t nc,
                                      int32_t w, int32_t out_height,
                                      int32_t out_width) {
  auto j = std::int32_t(blockIdx.x * blockDim.x + threadIdx.x);  // dz
  auto k = std::int32_t(blockIdx.y * blockDim.y + threadIdx.y);  // bs
  if ((j < dz) && (k < bs)) {
    for (int32_t n{0}; n < nc * w; ++n) {
      auto adr = k * dz * nc * w + n * dz + j;
      if (circle_map[k * dz * out_height * out_width +
                     j * out_height * out_width + idx_ptr[adr]] < 1)
        circle_map[k * dz * out_height * out_width +
                   j * out_height * out_width + idx_ptr[adr]] = 1.0;
    }
  }
}

__global__ void get_ray_map_kernel(float* ray_map, const int32_t* idx_ptr,
                                   int32_t bs, int32_t dz, int32_t nc,
                                   int32_t w, int32_t out_height,
                                   int32_t out_width) {
  auto j = std::int32_t(blockIdx.x * blockDim.x + threadIdx.x);  //  nc*w
  auto i = std::int32_t(blockIdx.y * blockDim.y + threadIdx.y);  // bs
  if ((j < nc * w) && (i < bs)) {
    auto adr = i * nc * w * dz + j * dz;
    for (int32_t n{0}; n < dz; ++n) {
      if (ray_map[i * nc * w * out_height * out_width +
                  j * out_height * out_width + idx_ptr[adr + n]] < 1)
        ray_map[i * nc * w * out_height * out_width +
                j * out_height * out_width + idx_ptr[adr + n]] = 1;
    }
  }
}

__global__ void clean_ray_map_value(float* ray_map, int32_t index, int32_t bs,
                                    int32_t nc, int32_t w, int32_t out_height,
                                    int32_t out_width) {
  auto idx = std::int32_t(blockIdx.x * blockDim.x + threadIdx.x);
  if (idx < bs * nc * w) {
    const int cw = idx % (nc * w);
    const int b = idx / (nc * w);
    ray_map[b * nc * w * out_height * out_width + cw * out_height * out_width +
            index] = 0;
  }
}

__global__ void clean_circle_map_value(float* circle_map, int32_t index,
                                       int32_t bs, int32_t dz,
                                       int32_t out_height, int32_t out_width) {
  auto idx = std::int32_t(blockIdx.x * blockDim.x + threadIdx.x);
  if (idx < bs * dz) {
    const int d = idx % dz;
    const int b = idx / dz;
    circle_map[b * dz * out_height * out_width + d * out_height * out_width +
               index] = 0;
  }
}

__device__ __forceinline__ static void reduceMax(float* address, float val) {
  int* address_as_i = reinterpret_cast<int*>(address);
  int old = *address_as_i, assumed;
  do {
    assumed = old;
    old = atomicCAS(address_as_i, assumed,
                    __float_as_int(fmaxf(val, __int_as_float(assumed))));
  } while (assumed != old || __int_as_float(old) < val);
}

__global__ void circle_map_max_and_clip_value(
    const float* circle_map, float* circle_map_max_clip, int32_t bs, int32_t dz,
    int32_t out_height, int32_t out_width, float clip_min) {
  auto idx = std::int32_t(blockIdx.x * blockDim.x + threadIdx.x);
  // data format:[bs, dz, out_height*out_width]
  if (idx < bs * dz * out_height * out_width) {
    const int c = idx % (out_height * out_width);
    const int d = idx / (out_height * out_width) % dz;
    const int b = idx / (out_height * out_width) / dz;
    // clip
    float inputElement = circle_map[b * dz * out_height * out_width +
                                    d * out_height * out_width + +c];
    float tmp = inputElement > clip_min ? inputElement : clip_min;
    // max
    reduceMax(&circle_map_max_clip[b * out_height * out_width + c], tmp);
  }
}

__global__ void ray_map_max_and_clip_value(const float* ray_map,
                                           float* ray_map_max_clip, int32_t bs,
                                           int32_t ncw, int32_t out_height,
                                           int32_t out_width, float clip_min) {
  auto idx = std::int32_t(blockIdx.x * blockDim.x + threadIdx.x);
  // data format:[bs, ncw, out_height*out_width]
  if (idx < bs * ncw * out_height * out_width) {
    const int c = idx % (out_height * out_width);
    const int n = idx / (out_height * out_width) % ncw;
    const int b = idx / (out_height * out_width) / ncw;
    // clip
    float inputElement = ray_map[b * ncw * out_height * out_width +
                                 n * out_height * out_width + +c];
    float tmp = inputElement > clip_min ? inputElement : clip_min;
    // max
    reduceMax(&ray_map_max_clip[b * out_height * out_width + c], tmp);
  }
}

__global__ void circle_map_div(float* circle_map,
                               const float* circle_map_max_clip, int32_t bs,
                               int32_t dz, int32_t out_height,
                               int32_t out_width) {
  auto idx = std::int32_t(blockIdx.x * blockDim.x + threadIdx.x);
  // data format:[bs, dz, out_height*out_width]    [bs, out_height*out_width]
  if (idx < bs * dz * out_height * out_width) {
    const int c = idx % (out_height * out_width);
    const int d = idx / (out_height * out_width) % dz;
    const int b = idx / (out_height * out_width) / dz;
    circle_map[b * dz * out_height * out_width + d * out_height * out_width +
               +c] /= circle_map_max_clip[b * out_height * out_width + c];
  }
}

__global__ void ray_map_div(float* ray_map, const float* ray_map_max_clip,
                            int32_t bs, int32_t ncw, int32_t out_height,
                            int32_t out_width) {
  auto idx = std::int32_t(blockIdx.x * blockDim.x + threadIdx.x);
  // data format:[bs, ncw, out_height*out_width]
  if (idx < bs * ncw * out_height * out_width) {
    const int c = idx % (out_height * out_width);
    const int n = idx / (out_height * out_width) % ncw;
    const int b = idx / (out_height * out_width) / ncw;
    ray_map[b * ncw * out_height * out_width + n * out_height * out_width +
            +c] /= ray_map_max_clip[b * out_height * out_width + c];
  }
}

namespace circle_ray {
extern "C" int32_t cal_coor_gpu(custom_coor_point* key_value, float* aug_mat,
                                float* ego_mat, int32_t sd_x, int32_t sd_y,
                                int32_t n_cam, int32_t dz, int32_t width,
                                int32_t height, std::vector<float> voxel_size,
                                std::vector<float> coors_range) {
  const float voxel_x = voxel_size[0];
  const float voxel_y = voxel_size[1];
  const float voxel_z = voxel_size[2];
  const float coors_x_min = coors_range[0];
  const float coors_y_min = coors_range[1];
  const float coors_z_min = coors_range[2];

  auto par_x = static_cast<float>((sd_x * width - 1.0) / (width - 1.0));
  auto par_y = static_cast<float>((sd_y * height - 1.0) / (height - 1.0));

  dim3 blocks_xyz(int32_t((width + 7) / 8), int32_t((height + 7) / 8),
                  int32_t(n_cam * dz + 15) / 16);
  dim3 threads_xyz(8, 8, 16);
  circle_ray_cal_idx_kernel<<<blocks_xyz, threads_xyz>>>(
      key_value, aug_mat, ego_mat, n_cam, dz, width, height, voxel_x, voxel_y,
      voxel_z, coors_x_min, coors_y_min, coors_z_min, par_x, par_y);
  cudaDeviceSynchronize();
  CHECK(cudaGetLastError());
  return 0;
}
extern "C" int32_t coor_reduce_hdim(custom_coor_point* o_coor_value,
                                    custom_coor_point* i_coor_value, int32_t bs,
                                    int32_t nc, int32_t dz, int32_t height,
                                    int32_t width) {
  dim3 blocks_xy(int32_t((width + 15) / 16), int32_t(bs * nc * dz + 63) / 64);
  dim3 threads_xy(16, 64);
  circle_ray_reduce_hdim<<<blocks_xy, threads_xy>>>(o_coor_value, i_coor_value,
                                                    bs, nc, dz, height, width);
  cudaDeviceSynchronize();
  CHECK(cudaGetLastError());
  return 0;
}
extern "C" int32_t coor_permute(custom_coor_point* o_coor_value,
                                custom_coor_point* i_coor_value,
                                int32_t* idx_ptr, int32_t bs, int32_t nc,
                                int32_t dz, int32_t width, int32_t out_height,
                                int32_t out_width) {
  const int32_t nthreads = bs * nc * dz * width;
  const int BS = 512;
  const int GS = (nthreads + BS - 1) / BS;
  permuteData_kernel<<<GS, BS>>>(o_coor_value, i_coor_value, idx_ptr, nthreads,
                                 bs, nc, dz, width, out_height, out_width);
  cudaDeviceSynchronize();
  CHECK(cudaGetLastError());
  return 0;
}
extern "C" int32_t coor_map(float* circle_map, float* ray_map,
                            const int32_t* idx_ptr, int32_t bs, int32_t nc,
                            int32_t dz, int32_t width, int32_t out_height,
                            int32_t out_width) {
  cudaStream_t stream_circle;
  cudaStream_t stream_ray;
  cudaStreamCreate(&stream_ray);
  cudaStreamCreate(&stream_circle);

  dim3 blocks_circle(int32_t((dz + 63) / 64), int32_t(bs + 1) / 2);
  dim3 threads_cirle(64, 2);
  get_circle_map_kernel<<<blocks_circle, threads_cirle, 0, stream_circle>>>(
      circle_map, idx_ptr, bs, dz, nc, width, out_height, out_width);

  dim3 blocks_ray(int32_t((nc * width + 63) / 64), int32_t(bs + 1) / 2);
  dim3 threads_ray(64, 2);
  get_ray_map_kernel<<<blocks_ray, threads_ray, 0, stream_ray>>>(
      ray_map, idx_ptr, bs, dz, nc, width, out_height, out_width);
  cudaDeviceSynchronize();
  CHECK(cudaGetLastError());

  const int BS = 512;
  const int GS_ray = (bs * nc * width + BS - 1) / BS;
  const int GS_circle = (bs * dz + BS - 1) / BS;

  const int invalid_idx = (out_height / 2) * out_width + out_width / 2;
  clean_ray_map_value<<<GS_ray, BS, 0, stream_ray>>>(
      ray_map, invalid_idx, bs, nc, width, out_height, out_width);
  clean_circle_map_value<<<GS_circle, BS, 0, stream_circle>>>(
      circle_map, invalid_idx, bs, dz, out_height, out_width);
  cudaDeviceSynchronize();
  CHECK(cudaGetLastError());

  //  // 归一化   max&&clip
  //  const int GS_ray_max_clip = (bs * nc * width * out_height * out_width+ BS
  //  - 1) / BS; const int GS_circle_max_clip = (bs * dz * out_height *
  //  out_width + BS - 1) / BS;
  //  circle_map_max_and_clip_value<<<GS_circle_max_clip, BS, 0,
  //  stream_circle>>>(circle_map, circle_map_max_clip, bs, dz, out_height,
  //  out_width, 1.0); ray_map_max_and_clip_value<<<GS_ray_max_clip, BS, 0,
  //  stream_ray>>>(ray_map, ray_map_max_clip, bs, nc*width, out_height,
  //  out_width, 1.0);
  //  // div
  //  circle_map_div<<<GS_circle_max_clip, BS, 0, stream_circle>>>(circle_map,
  //  circle_map_max_clip, bs, dz, out_height, out_width);
  //  ray_map_div<<<GS_ray_max_clip, BS, 0, stream_ray>>>(ray_map,
  //  ray_map_max_clip, bs, nc*width, out_height, out_width);
  //  cudaDeviceSynchronize();
  //  CHECK(cudaGetLastError());
  cudaStreamDestroy(stream_ray);
  cudaStreamDestroy(stream_circle);
}
}  // namespace circle_ray
